#!/usr/bin/python
"""
    qgram - a program to browse and analyze earthquake data 
    using spectrograms 
    Copyright (C) 2011 John Ivey

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>

Contact the authors through the Google Code website http://code.google.com/p/qgram

"""

import time
import calendar
import os
import ConfigParser

import Tkinter
import tkSimpleDialog
import tkFileDialog
import tkMessageBox

import numpy
import matplotlib
matplotlib.use('TkAgg')
import pylab
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2TkAgg
from matplotlib.figure import Figure

# qgram imports
import readers
import arrtimes

class App:
    def __init__(self,parent):
        menubar = Tkinter.Menu(parent)
        # Create file menu and add items
        filemenu = Tkinter.Menu(menubar, tearoff=0)
        filemenu.add_command(label="Import PSN type 4", command=self.import_psn)
        filemenu.add_command(label="Import delimited text", command=self.import_delimited_text)
        filemenu.add_separator()
        filemenu.add_command(label="Exit", command=self.on_exit)
        
        # Add file menu to menu bar
        menubar.add_cascade(label="File", menu=filemenu)
        
        #Add discrete commands to menu bar
        menubar.add_command(label = "Set NFFT", command = self.set_nfft)
        menubar.add_command(label = "Set cliphi", command = self.set_cliphi)
        menubar.add_command(label = "Set cliplo", command = self.set_cliplo)
        menubar.add_command(label = "Set autoclip", command = self.set_autoclip)
        menubar.add_command(label = "Show USGS arrival times", command = self.show_arrival)
        menubar.add_command(label = "test", command = self.test)
        
        #Add menubar to root window
        parent.config(menu=menubar)
        
        #Create status bar and add to bottom of root window
        self.status = StatusBar(parent)
        self.status.pack(side=Tkinter.BOTTOM, fill=Tkinter.X)
        
        self.fig = Figure(figsize=(5,4), dpi=100)
        
        self.canvas = FigureCanvasTkAgg(self.fig, master=root)
        self.canvas.get_tk_widget().pack(side=Tkinter.TOP, fill=Tkinter.BOTH, expand=1)

        self.canvas.mpl_connect('button_press_event', self.onclick)

        toolbar = NavigationToolbar2TkAgg( self.canvas, root )
        toolbar.update()
        self.canvas.get_tk_widget().pack(side=Tkinter.BOTTOM, fill=Tkinter.BOTH, expand=1)

        #self.subplot = self.fig.add_subplot(1,1,1)
        self.subplot = self.fig.add_axes([0.1, 0.1, 0.75, 0.8])
        self.bar = self.fig.add_axes([0.9, 0.1, 0.02, 0.8])
        self.subplot.set_title('')
        self.subplot.set_xlabel('Time (seconds)')
        self.subplot.set_ylabel('Frequency (Hz)')
        
        config = ConfigParser.SafeConfigParser()
        config.read(["config", os.path.expanduser("~/config")])
        section = config.sections()[0]
        items = config.items(section)

        configd = {}
        for key, value in items:
            configd[key] = value

        self.fileopenpath = "~/."
        if "fileopenpath" in configd:
            self.fileopenpath = configd["fileopenpath"]

        self.text_column = 0
        if "text_column" in configd:
            self.text_column = int(configd["text_column"])

        self.decimate = 1
        if "decimate" in configd:
            self.decimate = int(configd["decimate"])

        self.text_skiprows = 0
        if "text_skiprows" in configd:
            self.text_skiprows = float(configd["text_skiprows"])
        

        self.sample_frequency = 1
        if "sample_frequency" in configd:
            self.sample_frequency = float(configd["sample_frequency"])

        self.specgram_nfft = 2048
        if "specgram_nfft" in configd:
            self.specgram_nfft = int(configd["specgram_nfft"])
            
        self.text_delimiter = " "
        if "text_delimiter" in configd:
            self.text_delimiter = configd["text_delimiter"].split('"')[1]
        
        self.cliplo = 0.0
        self.cliphi = 76.0
        self.autoclip = True
        self.axis_dirty = True
        self.x = 0
        self.samples=0
        self.starttime = None

    def onclick(self, event):
        for i in range(0,len(self.bins)):
            if self.bins[i] > event.xdata:
                break
        for j in range(0,len(self.freqs)):
            if  self.freqs[j] > event.ydata:
                break
        val = self.Z[len(self.freqs)-j, i]
        print 'button=%d, x=%d, y=%d, xdata=%f, ydata=%f, i=%d, j=%d, val=%f' %(event.button, event.x, event.y, event.xdata, event.ydata, i, j, val)
        if i == 0:
            t = 0.0
        else:
            t = self.bins[i-1]
        self.status.set_status("NFFT:%6d    Sample rate:   %6.2f Hz    t: %6.2f sec   f: %6.3f Hz    val: %6.1f dB", self.specgram_nfft, self.sample_frequency, t, self.freqs[j-1], val)
        
        
        
        


    def update_specgram(self, recalculate=True):
        myaxis = self.subplot.axis()
        if recalculate == True:
            self.subplot.set_title('')
            self.subplot.set_xlabel('Time (s)')
            self.subplot.set_ylabel('Frequency (Hz)')
            self.Pxx, self.freqs, self.bins, im = self.subplot.specgram(self.x, 
                                                                        self.specgram_nfft, 
                                                                        self.sample_frequency, 
                                                                        noverlap = self.specgram_nfft/2, 
                                                                        detrend=pylab.detrend_linear, 
                                                                        interpolation="nearest")
        self.Z = 20. * numpy.log10(self.Pxx)
        self.Z = numpy.maximum(self.Z,-300.0)
        self.Z = numpy.flipud(self.Z)
        if self.autoclip == True:
            self.cliphi = numpy.amax(self.Z)
            self.cliplo = numpy.amin(self.Z)
        xextent = 0, numpy.amax(self.bins)
        xmin, xmax = xextent
        extent = xmin, xmax, self.freqs[0], self.freqs[-1]
        self.subplot.cla()
        im = self.subplot.imshow(self.Z, 
                                 cmap=None, 
                                 norm=None, 
                                 vmin = self.cliplo, 
                                 vmax = self.cliphi, 
                                 extent=extent, 
                                 interpolation = "nearest")
        self.subplot.axis('auto')
        self.canvas.draw()
        if self.axis_dirty == True:
            self.axis_dirty = False
        else:
            self.subplot.axis(myaxis)

        self.fig.colorbar(im, self.bar, extend = "both") # draw colorbar

        self.status.set_status("NFFT:%6d  Sample rate: %6.2f Hz", self.specgram_nfft, self.sample_frequency)
        self.canvas.draw()

    def import_psn(self):
        filename = tkFileDialog.askopenfilename( filetypes=(('psn File', '*.psn'),('All Files', '*')),
                                                 initialdir=os.path.dirname(self.fileopenpath),
                                                 initialfile=os.path.basename(self.fileopenpath))
        if filename:
            retval = tkSimpleDialog.askinteger("decimate", "decimation factor", initialvalue = self.decimate)
            if retval != None:
                self.decimate = retval
                self.x, self.sample_frequency = readers.read_psn(filename)
                # Need to get start time from PSN file
                if self.decimate != 1:
                    self.x = readers.decimate_data(self.x, self.decimate, 1000, "fir")
                    self.sample_frequency = self.sample_frequency / self.decimate
                self.axis_dirty = True
                self.update_specgram(recalculate = True)
        return

    def import_delimited_text(self):
        filename = tkFileDialog.askopenfilename( filetypes=(('txt File', '*.txt'),
                                                            ('All Files', '*')),
                                                 initialdir=os.path.dirname(self.fileopenpath),
                                                 initialfile=os.path.basename(self.fileopenpath))

        if filename:
            self.fileopenpath = filename;
            retfreq = tkSimpleDialog.askfloat("frequency", "enter sample frequency", initialvalue = self.sample_frequency)
            if retfreq != None:
                retdecimate = tkSimpleDialog.askinteger("decimate", 
                                                        "decimation factor", 
                                                        initialvalue = self.decimate)
                if retdecimate != None:
                    retcolumn = tkSimpleDialog.askinteger("column to read", 
                                                          "column number (first column=0)", 
                                                          initialvalue = self.text_column)
                    if retcolumn != None:
                        self.text_column = retcolumn
                        self.decimate = retdecimate
                        self.sample_frequency = retfreq
                        self.x = readers.read_text(filename, column = self.text_column, skiprows = self.text_skiprows, delim = self.text_delimiter)
                        if self.decimate != 1:
                            self.x = readers.decimate_data(self.x, self.decimate, 1000, "fir")
                            self.sample_frequency = self.sample_frequency/self.decimate
                        self.update_specgram(recalculate = True)
                        root.wm_title("qgram - "+filename)
                        self.starttime = None
        return

    def set_nfft(self):
#        myaxis = self.subplot.axis()
        retval = tkSimpleDialog.askinteger("NFFT", "Set length of transform", initialvalue = self.specgram_nfft)
        if retval==None or retval>numpy.alen(self.x) or retval<=128:
            return
        self.specgram_nfft = retval
        self.update_specgram(recalculate = True)
#        self.subplot.axis('auto')
#        self.canvas.draw()

    def set_cliphi(self):
        retval = tkSimpleDialog.askfloat("High clip (dB)", "set high limit of image (red)", initialvalue = self.cliphi)
        if retval==None:
            return
        self.cliphi = retval
        self.autoclip = False
        self.update_specgram(recalculate = False)

    def set_cliplo(self):
#        myaxis = self.subplot.axis()
        retval = tkSimpleDialog.askfloat("Low clip (dB)", "set low limit of image (blue)", initialvalue = self.cliplo)
        if retval != None:
            self.cliplo = retval
            self.autoclip = False
            self.update_specgram(recalculate = False)
#            self.subplot.axis(myaxis)
#            self.canvas.draw()

    def on_exit(self):
        root.destroy()

    def set_autoclip(self):
        self.autoclip = True
        self.canvas.draw()
        
    def test(self):
        self.x = numpy.sin(numpy.arange(0,1000000*6.28,6.28/10))
        self.update_specgram()
        self.canvas.draw()

    def show_arrival(self):
        if self.starttime == None:
            retval1 = None
            initstring = ""
            while retval1 == None:
                retval = tkSimpleDialog.askstring("File start date and time", 
                                                  "Date and time 'yyyy-mm-dd hhmmss'", 
                                                  initialvalue = initstring)
                initstring = retval1
                if retval != None:
                    try:
                        retval1 = time.strptime(retval, "%Y-%m-%d %H%M%S")
                    except ValueError as inst:
                        tkMessageBox.showwarning("Error", inst)
                    else:
                        self.starttime = retval1
                
        tstart = calendar.timegm(self.starttime)
        tend = tstart + numpy.alen(self.x)
        quakes = arrtimes.getquakes()
        offset = 0
        for quake in quakes:
            for phase in quake["phases"]:
                t = calendar.timegm(phase["arrtime"])
                if t<tend and t> tstart:
                    if phase["phase"] == "LQ":
                        a=self.subplot.axes.annotate(quake["comment"]+" "+str(quake["mag"]), 
                                                     xy=(t-tstart, .1),
                                                     xycoords='data',
                                                     xytext=(-50, 30+offset),
                                                     textcoords='offset points',
                                                     arrowprops=dict(arrowstyle="->",
                                                                     connectionstyle="angle,angleA=0,angleB=-90,rad=10"))
                        offset = offset + 10
                    if phase["phase"] == "P":
                        a=self.subplot.axes.annotate("", xy=(t-tstart, .1),  xycoords='data',
                        xytext=(0, 30+offset), textcoords='offset points',
                        arrowprops=dict(arrowstyle="->",connectionstyle="arc3,rad=0")
                        )        
                    if phase["phase"] == "S":
                        a=self.subplot.axes.annotate("", xy=(t-tstart, .1),  xycoords='data',
                        xytext=(0, 30+offset), textcoords='offset points',
                        arrowprops=dict(arrowstyle="->",connectionstyle="arc3,rad=0")
                        )        
        self.canvas.draw()

class StatusBar(Tkinter.Frame):
    def __init__(self, master):
        Tkinter.Frame.__init__(self, master)
        self.label = Tkinter.Label(self, bd=1, relief=Tkinter.SUNKEN, anchor=Tkinter.W)
        self.label.pack(fill=Tkinter.X)

    def set_status(self, pfformat, *args):
        self.label.config(text=pfformat % args)
        self.label.update_idletasks()

    def clear_status(self):
        self.label.config(text="")
        self.label.update_idletasks()

root = Tkinter.Tk()
root.wm_title("qgram")
app = App(root)
root.protocol("WM_DELETE_WINDOW", app.on_exit)
root.mainloop()